/* -*-  Mode: C++; c-file-style: "gnu"; indent-tabs-mode:nil; -*- */
  //
// Copyright (c) 2006 Georgia Tech Research Corporation
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License version 2 as
// published by the Free Software Foundation;
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
//
// Author: George F. Riley<riley@ece.gatech.edu>
//

// ns3 - On/Off Data Source Application class
// George F. Riley, Georgia Tech, Spring 2007
// Adapted from ApplicationFlow in GTNetS.

#include "ns3/log.h"
#include "ns3/address.h"
#include "ns3/node.h"
#include "ns3/nstime.h"
#include "ns3/data-rate.h"
#include "ns3/random-variable.h"
#include "ns3/socket.h"
#include "ns3/simulator.h"
#include "ns3/socket-factory.h"
#include "ns3/packet.h"
#include "ns3/uinteger.h"
#include "ns3/inet-socket-address.h"
#include "ns3/trace-source-accessor.h"
#include "flow-application.h"
#include "ns3/tcp-socket-factory.h"
#include "ns3/tcp-header.h"
#include "ns3/ipv4-raw-socket-factory.h"

NS_LOG_COMPONENT_DEFINE ("FlowApplication");

using namespace std;

namespace ns3 {

NS_OBJECT_ENSURE_REGISTERED (FlowApplication);
  
TypeId
FlowApplication::GetTypeId (void)
{
  static TypeId tid = TypeId ("ns3::FlowApplication")
    .SetParent<Application> ()
    .AddConstructor<FlowApplication> ()
    .AddAttribute ("Remote", "The address of the destination",
                   AddressValue (),
                   MakeAddressAccessor (&FlowApplication::m_remote),
                   MakeAddressChecker ())
    .AddAttribute ("Local", "The Address on which to Bind the rx socket.",
                   AddressValue (),
                   MakeAddressAccessor (&FlowApplication::m_local),
                   MakeAddressChecker ())             
    .AddAttribute ("FlowNum", "Number of concurrent TCP flows",
                   UintegerValue (0),
                   MakeUintegerAccessor (&FlowApplication::m_flowNum),
                   MakeUintegerChecker<uint32_t>())
    .AddAttribute ("FlowSize", "A RandomVariable which sets flow size in bytes",
                   RandomVariableValue (ConstantVariable (0)),
                   MakeRandomVariableAccessor (&FlowApplication::m_flowSize),
                   MakeRandomVariableChecker ())                                
    .AddAttribute ("StatCallback", "A callback for logging flow duration",
                   CallbackValue (),
                   MakeCallbackAccessor (&FlowApplication::m_statCallback),
                   MakeCallbackChecker ())
       ;
  return tid;
}

TypeId
Flow::GetTypeId ()
{
  static TypeId tid = TypeId("ns3::Flow")
    .SetParent<Object> ()
    ;
  return tid;
}
    
TypeId
FlowSink::GetTypeId ()
{
  static TypeId tid = TypeId("ns3::FlowSink")
    .SetParent<Object> ()
    ;
  return tid;
}
       
Flow::Flow ()
{}      

Flow::~Flow ()
{}

void
Flow::SetFinishCallback( Callback<void, Ptr<Flow> > flowFinished )
{
  //NS_LOG_FUNCTION_NOARGS ();
  m_flowFinished = flowFinished;
}

void
Flow::Start (uint32_t size)
{
  NS_ASSERT ( size );
  m_flowSize   = size;
  m_dataToSend = size;

  m_socket = Socket::CreateSocket (m_node, TcpSocketFactory::GetTypeId());
  m_socket->Bind ();
  m_txMax = m_socket->GetTxAvailable();

  m_startTime = Simulator::Now();
  m_socket->Connect (m_remote);                  
  m_socket->SetDataSentCallback(MakeCallback(&Flow::HandleSend, this));

  HandleSend(m_socket,0);
}                                                      
 
void
Flow::Restart (uint32_t size) 
{
  // default to previously set size
  if ( size )
    m_flowSize = size;

  Start(m_flowSize);
}
   
uint32_t
Flow::GetSize () 
{
  return m_flowSize;
}                   

Time
Flow::GetDuration () 
{
  return m_endTime - m_startTime;
}                   

void
Flow::HandleClose(Ptr<Socket> socket) 
{
  //NS_LOG_FUNCTION (this << socket << Simulator::Now());
  Ptr<Packet> packet;
  Address from;
  while (packet = socket->RecvFrom (from))
  {
    if (packet->GetSize() == 0)
    { //EOF
      m_endTime = Simulator::Now();
      socket->Close();
      socket->SetRecvCallback (MakeNullCallback<void, Ptr<Socket> > ());
      // notify caller
      if ( !m_flowFinished.IsNull() )
      {
        m_flowFinished (this);
      }
    }  
  }  
}

void 
Flow::SetNode(Ptr<Node> node)
{
  m_node = node;
}
void 
Flow::SetRemote(const Address &peer)
{
  if(!InetSocketAddress::IsMatchingType (peer))
  {
    // fail silently for now
    return;
  }
  m_remote = peer;
}

void 
Flow::End ()
{
  // close connection, cancel callback
  NS_LOG_FUNCTION(this << m_socket);
  m_socket->SetRecvCallback( MakeNullCallback<void, Ptr<Socket> > ());
  m_socket->Close();
  m_socket = 0;
}

void 
Flow::HandleSend(Ptr<Socket> socket, uint32_t dataSent) 
{
  if ( !m_dataToSend )
  {
    //NS_LOG_LOGIC ("Waiting for remote peer to close connection " << socket);
    m_socket->SetDataSentCallback (MakeNullCallback<void, Ptr<Socket>,uint32_t > ());
    m_socket->SetRecvCallback(MakeCallback(&Flow::HandleClose, this));
    return;
  }

  if ( dataSent < m_txMax - m_socket->GetTxAvailable() )
    return;

  //NS_LOG_LOGIC ("Sending data " << dataSent <<  " " <<m_dataToSend << " " << m_flowSize);
  uint32_t dataToSend = min ( m_dataToSend , m_socket->GetTxAvailable() );
  
  //send size in first packet
  Ptr<Packet> packet;
  if ( m_flowSize == m_dataToSend )
  {
    uint8_t buffer[dataToSend];
    memset(buffer, 0, dataToSend);

    for( uint8_t i = 0; i < 4 ; i++ )
    {
      buffer[i] = (( m_flowSize >> i*8) & 0xff );
    }
 
    packet = Create<Packet> (buffer, dataToSend);
  }
  else
  {
    packet = Create<Packet> (dataToSend);
  }
  m_socket->Send (packet);
  m_dataToSend -= dataToSend;
}
 
FlowSink::FlowSink () :
m_dataToRecv(0)
{
}                 
 
void 
FlowSink::End ()
{
  // do not close socket, terminate it silently
  // FIN could be misinterpreted as successful flow finish
  m_socket->Close();
}      

void 
FlowSink::HandleRead(Ptr<Socket> socket)
{
  Ptr<Packet> packet;
  Address from;

  while (packet = socket->RecvFrom (from))
    {
      if ( !m_dataToRecv )
      {
        // retrieve size from first packet
        // TODO: could use packet tag? 
        uint8_t buffer[packet->GetSize()];
        packet->CopyData(buffer, packet->GetSize());

        for ( int8_t i = 3; i >= 0 ; i-- )
        {
          m_dataToRecv <<= 8;
          m_dataToRecv |= buffer[i];
        }
      }

      m_dataToRecv -= packet->GetSize();
      if ( !m_dataToRecv )
      {
        //NS_LOG_LOGIC ("Tear connection");
        socket->Close();
      }

    }    
}
          
void
FlowSink::SetSocket ( Ptr<Socket> socket )
{          
  if (m_socket)
  {
    NS_ASSERT(true);
  }
  m_socket = socket;
  m_socket->SetRecvCallback (MakeCallback(&FlowSink::HandleRead, this));
}
          
FlowApplication::FlowApplication ()
{
  NS_LOG_FUNCTION_NOARGS ();
  m_statCallback = MakeNullCallback<void, uint32_t, Time > ();
}

FlowApplication::~FlowApplication()
{
  NS_LOG_FUNCTION_NOARGS ();
}


void
FlowApplication::DoDispose (void)
{
  NS_LOG_FUNCTION_NOARGS ();

  // chain up
  Application::DoDispose ();
}

// Application Methods
void 
FlowApplication::StartApplication() // Called at time specified by Start
{
  NS_LOG_FUNCTION (this);

  // sink 
  if ( !m_socket && !m_local.IsInvalid() )
  {
    //NS_LOG_LOGIC ("creating listening socket " << m_local);
    m_socket = Socket::CreateSocket (GetNode(), TcpSocketFactory::GetTypeId());
    m_socket->Bind (m_local);
    m_socket->Listen ();

    m_socket->SetAcceptCallback (
            MakeNullCallback<bool, Ptr<Socket>, const Address &> (),
            MakeCallback(&FlowApplication::HandleAccept, this));
  }

  // source 
  for ( unsigned int i = 0; i < m_flowNum; i++ )
  {
    Ptr<Flow> flow = CreateObject<Flow> ();
    flow->SetNode(GetNode());
    flow->SetRemote(m_remote);
    flow->SetFinishCallback(MakeCallback(&FlowApplication::FlowFinishedCallback, this));
                       
    uint32_t size;
    do {
      size = (uint32_t)m_flowSize.GetValue();
    } while (!size );
                       
    flow->Start(size);
    m_flowList.push_back(flow);
    //NS_LOG_FUNCTION("Created flow " << flow );
  }
}    

void 
FlowApplication::FlowFinishedCallback(Ptr<Flow> flow)
{
  if(flow)
  {
    if (!m_statCallback.IsNull ())
    {
      m_statCallback ( flow->GetSize() , flow->GetDuration() );
    }

    uint32_t size;
    do {
      size = (uint32_t)m_flowSize.GetValue();
    } while (!size );

    flow->Restart(size);
  }
}    

void 
FlowApplication::StopApplication() // Called at time specified by Stop
{
  NS_LOG_FUNCTION_NOARGS ();

  //for sink, stop accepting new flows
  if( m_socket )
  {
    m_socket->Close();
    m_socket = 0;
  }                                
  
  // ...and opened sockets
  while(!m_sinkList.empty()) 
  {
    Ptr<FlowSink> sink = m_sinkList.front(); 
    sink->End();
    m_sinkList.pop_front();
  }  

  //shut sources down.
  while(!m_flowList.empty()) 
  {
    Ptr<Flow> flow = m_flowList.front(); 
    flow->End();
    m_flowList.pop_front();
  }
 

}


void 
FlowApplication::HandleAccept (Ptr<Socket> s, const Address& from)
{
  //NS_LOG_FUNCTION (this << s << from);
  //create Flow sink
  Ptr<FlowSink> sink = CreateObject<FlowSink> ();
  sink->SetSocket(s);
  m_sinkList.push_back (sink);
}


} // Namespace ns3
